import streamlit as st
import asyncio
import config
import db
from crawler_manager import CrawlerManager
from tools import utils
import nest_asyncio
from webui_db import DataManager

# 启用嵌套事件循环支持
nest_asyncio.apply()

# 设置页面配置
st.set_page_config(
    page_title="社交媒体数据采集平台",
    page_icon="🔍",
    layout="wide"
)

# 全局变量
platform_map = {"抖音": "dy", "微博": "wb", "小红书": "xhs"}

# 初始化数据库连接和事件循环
if 'db_initialized' not in st.session_state:
    st.session_state.db_initialized = False
    st.session_state.data_manager = None
    st.session_state.loop = None

if not st.session_state.db_initialized:
    try:
        st.session_state.db_initialized = True
        st.session_state.loop = asyncio.new_event_loop()
        asyncio.set_event_loop(st.session_state.loop)
        st.session_state.loop.run_until_complete(db.init_db())
        st.session_state.data_manager = DataManager()
        st.success("数据库连接成功！")
    except Exception as e:
        st.error(f"数据库连接失败：{str(e)}")

# 主界面
st.title("社交媒体数据采集平台")

# 平台选择
platform = st.selectbox(
    "选择平台",
    ["抖音", "微博", "小红书"]
)

# 数据采集场景选择
scenario = st.selectbox(
    "选择数据采集场景",
    [
        "场景1：关键词搜索数据采集",
        "场景2：指定用户数据采集",
        "场景3：用户粉丝数据采集"
    ]
)

def get_user_list_sync(platform_code: str) -> list:
    """同步方式获取用户列表"""
    try:
        if not st.session_state.data_manager or not st.session_state.loop:
            return []
            
        async def _get_users():
            try:
                # 使用空字符串作为关键词，获取所有用户
                users = await st.session_state.data_manager.get_unique_user_ids_and_sec_uids(
                    platform=platform_code,
                    keyword="",
                    limit=100
                )
                return users
            except Exception as e:
                utils.logger.error(f"获取用户列表失败: {str(e)}")
                return []
                
        return st.session_state.loop.run_until_complete(_get_users())
    except Exception as e:
        utils.logger.error(f"同步获取用户列表失败: {str(e)}")
        return []

# 任务配置区域
with st.expander("任务配置", expanded=True):
    col1, col2 = st.columns(2)
    
    with col1:
        if scenario == "场景1：关键词搜索数据采集":
            keywords = st.text_area(
                "输入搜索关键词（多个关键词用逗号分隔）",
                placeholder="例如：比亚迪,新能源,汽车"
            )
            base_limit = st.number_input("每个关键词获取的帖子数量", min_value=1, value=10)
            comments_limit = st.number_input("每个帖子获取的评论数量", min_value=1, value=20)
            
        elif scenario == "场景2：指定用户数据采集":
            # 获取用户列表
            platform_code = platform_map.get(platform)
            if platform_code:
                user_list = get_user_list_sync(platform_code)
                
                if user_list:
                    # 创建用户选择列表，包含更多信息
                    user_options = []
                    user_labels = {}
                    for user_id, nickname in user_list:
                        option = f"{user_id}"
                        user_options.append(option)
                        user_labels[option] = f"{nickname} (ID: {user_id})"
                    
                    # 多选框
                    selected_users = st.multiselect(
                        "选择用户（可多选）",
                        options=user_options,
                        format_func=lambda x: user_labels[x],
                        help="从列表中选择要采集数据的用户"
                    )
                    
                    if selected_users:
                        user_ids = "\n".join(selected_users)
                    else:
                        user_ids = st.text_area(
                            "或手动输入用户ID（每行一个）",
                            placeholder="例如：\nuser1\nuser2"
                        )
                else:
                    st.warning("未能获取到用户列表，请手动输入用户ID")
                    user_ids = st.text_area(
                        "输入用户ID（每行一个）",
                        placeholder="例如：\nuser1\nuser2"
                    )
                    
            posts_limit = st.number_input("每个用户获取的帖子数量", min_value=1, value=20)
            comments_limit = st.number_input("每个帖子获取的评论数量", min_value=1, value=50)
            
        else:  # 场景3：用户粉丝数据采集
            # 获取用户列表
            platform_code = platform_map.get(platform)
            if platform_code:
                user_list = get_user_list_sync(platform_code)
                
                if user_list:
                    # 创建用户选择列表，包含更多信息
                    user_options = []
                    user_labels = {}
                    for user_id, nickname in user_list:
                        option = f"{user_id}"
                        user_options.append(option)
                        user_labels[option] = f"{nickname} (ID: {user_id})"
                    
                    # 多选框
                    selected_users = st.multiselect(
                        "选择用户（可多选）",
                        options=user_options,
                        format_func=lambda x: user_labels[x],
                        help="从列表中选择要采集粉丝数据的用户"
                    )
                    
                    if selected_users:
                        user_ids = "\n".join(selected_users)
                    else:
                        user_ids = st.text_area(
                            "或手动输入用户ID（每行一个）",
                            placeholder="例如：\nuser1\nuser2"
                        )
                else:
                    st.warning("未能获取到用户列表，请手动输入用户ID")
                    user_ids = st.text_area(
                        "输入用户ID（每行一个）",
                        placeholder="例如：\nuser1\nuser2"
                    )
                    
            base_limit = st.number_input("每个用户获取的粉丝数量", min_value=1, value=100)
            posts_limit = st.number_input("每个粉丝获取的帖子数量", min_value=1, value=20)
            comments_limit = st.number_input("每个帖子获取的评论数量", min_value=1, value=50)

# 创建状态显示组件
progress_bar = st.progress(0)
status_text = st.empty()

def update_status(message: str, progress: float):
    """更新状态显示"""
    status_text.text(message)
    progress_bar.progress(progress)

async def run_crawler_task():
    """执行爬虫任务"""
    async with CrawlerManager() as manager:
        try:
            utils.logger.info("[execute_task] 初始化爬虫管理器")
            
            if scenario == "场景1：关键词搜索数据采集":
                if not keywords:
                    st.error("请输入关键词！")
                    return
                    
                utils.logger.info(f"[execute_task] 关键词: {keywords}")
                await manager.collect_simple_search_data(
                    platform=platform,
                    keywords=keywords,
                    posts_limit=base_limit,
                    comments_limit=comments_limit,
                    status_callback=update_status
                )
                
            elif scenario == "场景2：指定用户数据采集":
                if not user_ids:
                    st.error("请输入或选择用户ID！")
                    return
                    
                user_id_list = [uid.strip() for uid in user_ids.split("\n") if uid.strip()]
                utils.logger.info(f"[execute_task] 用户IDs: {user_id_list}")
                await manager.collect_user_posts_data(
                    platform=platform,
                    user_ids=user_id_list,
                    posts_limit=posts_limit,
                    comments_limit=comments_limit,
                    status_callback=update_status
                )
                
            else:  # 场景3：用户粉丝数据采集
                if not user_ids:
                    st.error("请输入或选择用户ID！")
                    return
                    
                user_id_list = [uid.strip() for uid in user_ids.split("\n") if uid.strip()]
                utils.logger.info(f"[execute_task] 用户IDs: {user_id_list}")
                await manager.collect_fans_data(
                    platform=platform,
                    user_ids=user_id_list,
                    base_limit=base_limit,
                    posts_limit=posts_limit,
                    comments_limit=comments_limit,
                    status_callback=update_status
                )
            
            st.success("数据采集完成！")
            
        except Exception as e:
            error_msg = f"采集过程中出现错误: {str(e)}"
            utils.logger.error(f"[execute_task] {error_msg}")
            st.error(error_msg)

def execute_task():
    """执行任务的包装函数"""
    if not st.session_state.loop:
        st.error("事件循环未初始化")
        return
        
    try:
        st.session_state.loop.run_until_complete(run_crawler_task())
    except Exception as e:
        st.error(f"执行任务时出错: {str(e)}")

# 执行按钮
if st.button("开始采集", type="primary"):
    with st.spinner("正在执行数据采集..."):
        execute_task() 